import torch
from torch import nn
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.yolo import YOLOv8
import hiddenlayer as hl
from torchview import draw_graph
import argparse
from graphviz import Digraph
from torchviz import make_dot


def visualize_with_hiddenlayer(model, save_path):
    """使用hiddenlayer可视化模型"""
    transforms = [
        hl.transforms.Fold("Conv > BatchNorm > ReLU", "ConvBlock"),
        hl.transforms.Fold("ConvBlock > ConvBlock", "DoubleConvBlock"),
    ]
    graph = hl.build_graph(model, torch.zeros([1, 3, 640, 640]), transforms=transforms)
    graph.save(save_path, format='png')


def visualize_with_torchview(model, save_path):
    """使用torchview可视化模型"""
    model_graph = draw_graph(
        model, 
        input_size=(1, 3, 640, 640),
        expand_nested=True,
        hide_inner_tensors=True,
        hide_module_functions=True
    )
    model_graph.visual_graph.render(save_path, format='png', cleanup=True)


def visualize_with_torchviz(model, save_path):
    """使用torchviz可视化模型"""
    x = torch.zeros(1, 3, 640, 640)
    y = model(x)
    dot = make_dot(y, params=dict(model.named_parameters()))
    dot.render(save_path, format='png', cleanup=True)


def visualize_with_graphviz(model, save_path):
    """使用graphviz手动创建可视化图"""
    def add_nodes(graph, module, parent_name=''):
        for name, layer in module.named_children():
            node_name = f"{parent_name}/{name}" if parent_name else name
            # 添加节点
            graph.node(node_name, f'{name}\n{str(layer.__class__.__name__)}')
            # 如果有父节点，添加边
            if parent_name:
                graph.edge(parent_name, node_name)
            # 递归处理子模块
            add_nodes(graph, layer, node_name)

    dot = Digraph(comment='Model Architecture')
    dot.attr(rankdir='TB')  # 从上到下的布局
    add_nodes(dot, model)
    dot.render(save_path, format='png', cleanup=True)


def main():
    parser = argparse.ArgumentParser(description='Visualize YOLOv8 model architecture')
    parser.add_argument('--method', type=str, default='all',
                        choices=['hiddenlayer', 'torchview', 'torchviz', 'graphviz', 'all'],
                        help='Visualization method to use')
    parser.add_argument('--output-dir', type=str, default='model_plots',
                        help='Output directory for visualization files')
    args = parser.parse_args()

    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)

    # 创建模型实例
    model = YOLOv8(num_classes=80)
    
    # 设置为评估模式
    model.eval()

    methods = {
        'hiddenlayer': (visualize_with_hiddenlayer, 'hiddenlayer_model'),
        'torchview': (visualize_with_torchview, 'torchview_model'),
        'torchviz': (visualize_with_torchviz, 'torchviz_model'),
        'graphviz': (visualize_with_graphviz, 'graphviz_model')
    }

    if args.method == 'all':
        for method_name, (func, filename) in methods.items():
            try:
                print(f"\nGenerating visualization using {method_name}...")
                save_path = os.path.join(args.output_dir, filename)
                func(model, save_path)
                print(f"Saved visualization to {save_path}.png")
            except Exception as e:
                print(f"Error with {method_name}: {str(e)}")
    else:
        func, filename = methods[args.method]
        save_path = os.path.join(args.output_dir, filename)
        func(model, save_path)
        print(f"\nSaved visualization to {save_path}.png")


if __name__ == '__main__':
    main()
